package com.xzx.flink.tableapi;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.functions.AggregateFunction;

import static org.apache.flink.table.api.Expressions.$;

/**
 * 自定义UDF聚合函数
 *
 * @version 1.0
 * @auther xinzhixuan
 * @date 2022/6/11 21:58
 */
public class TableAPI_07_UDF_AggregateFunction {
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<Tuple2<Long, Integer>> stream = env.fromElements(Tuple2.of(1L,2), Tuple2.of(2L, 1));
        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);
        tableEnv.createTemporaryView("my_table", stream, Schema.newBuilder()
                .columnByExpression("num", $("f0"))
                .columnByExpression("weight", $("f1"))
                .build());

        tableEnv.createTemporarySystemFunction("weight_avg", WeightedAvgFunction.class);

        tableEnv.toChangelogStream(tableEnv.sqlQuery("select weight_avg(num, weight) from my_table")).print("-->");

        env.execute();

    }

    // 累加器类型定义
    public static class WeightedAvgAccumulator {
        public long sum = 0; // 加权和
        public int count = 0; // 数据个数
    }

    // 自定义聚合函数，输出为长整型的平均值
    public static class WeightedAvgFunction extends AggregateFunction<Double, WeightedAvgAccumulator> {


        @Override
        public Double getValue(WeightedAvgAccumulator accumulator) {
            if (accumulator.count == 0) {
                return 0.0;
            }
            return accumulator.sum * 1.0 / accumulator.count;
        }

        @Override
        public WeightedAvgAccumulator createAccumulator() {
            return new WeightedAvgAccumulator();
        }

        // 累加计算方法，每来一行数据都会调用, 这个方法也是底层约定好的
        public void accumulate(WeightedAvgAccumulator acc, Long iValue, Integer iWeight) {
            acc.sum += iValue * iWeight;
            acc.count += iWeight;
        }
    }
}


